AES

---

package encryption_utils

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"strings"
)

// AES struct to hold key and IV
type AES struct {
	key []byte
	iv  []byte
}

// Initialize AES with random key and IV
func NewAES() *AES {
	// Generate random key and IV
	key := make([]byte, 32)           // 32 bytes for AES-256
	iv := make([]byte, aes.BlockSize) // 16 bytes for AES block size

	if _, err := rand.Read(key); err != nil {
		panic(err)
	}
	if _, err := rand.Read(iv); err != nil {
		panic(err)
	}

	return &AES{
		key: key,
		iv:  iv,
	}
}

// Represents the response from AES encryption
type EncryptedResponse struct {
	EncryptedData string `json:"encryptedData"`
	AESProperties string `json:"aesProperties"`
	AES           bool   `json:"aes"`
}

// PKCS7 padding implementation
func pkcs7Pad(data []byte, blockSize int) []byte {
	padding := blockSize - len(data)%blockSize
	padtext := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(data, padtext...)
}

// PKCS7 unpadding implementation
func pkcs7Unpad(data []byte, blockSize int) ([]byte, error) {
	if len(data) == 0 || len(data)%blockSize != 0 {
		return nil, fmt.Errorf("invalid data length")
	}
	padLen := int(data[len(data)-1])
	if padLen == 0 || padLen > blockSize {
		return nil, fmt.Errorf("invalid padding length")
	}
	for i := len(data) - padLen; i < len(data); i++ {
		if data[i] != byte(padLen) {
			return nil, fmt.Errorf("invalid padding byte")
		}
	}
	return data[:len(data)-padLen], nil
}

func (a *AES) EncryptWithAES(data string, r *RSA) (*EncryptedResponse, error) {
	// Create AES cipher
	block, err := aes.NewCipher(a.key)
	if err != nil {
		return nil, err
	}

	// Pad data to be multiple of block size
	dataInBytes := []byte(data)
	padded := pkcs7Pad(dataInBytes, aes.BlockSize)

	// Encrypt data using CBC mode
	cbc := cipher.NewCBCEncrypter(block, a.iv)

	// Allocate space for ciphertext
	ciphertext := make([]byte, len(padded))
	// Encrypt the padded data
	cbc.CryptBlocks(ciphertext, padded)

	// Encode AES properties (key and IV) in base64 and concatenate
	ivBase64 := base64.StdEncoding.EncodeToString(a.iv)
	keyBase64 := base64.StdEncoding.EncodeToString(a.key)

	// Concatenate key and IV with a dot separator
	properties := fmt.Sprintf("%s.%s", keyBase64, ivBase64)

	// Encode ciphertext in base64
	encryptedData := base64.StdEncoding.EncodeToString(ciphertext)
	// Encrypt AES properties using RSA public key
	aesProperties, err := r.EncryptUsingPublicKey(properties)
	if err != nil {
		return nil, err
	}

	return &EncryptedResponse{
		EncryptedData: encryptedData,
		AESProperties: aesProperties,
		AES:           true,
	}, nil
}

func (a *AES) DecryptWithAES(encryptedAESProperties, encryptedData string, r *RSA) (string, error) {
	// Decrypt AES properties using RSA private key
	decryptedProperties, err := r.DecryptUsingPrivateKey(encryptedAESProperties)
	if err != nil {
		return "", err
	}

	// Split the decrypted properties to get key and IV
	parts := strings.Split(decryptedProperties, ".")
	if len(parts) != 2 {
		return "", fmt.Errorf("invalid AES properties format")
	}

	iv, err := base64.StdEncoding.DecodeString(parts[1])
	if err != nil {
		return "", fmt.Errorf("failed to decode IV: %v", err)
	}

	key, err := base64.StdEncoding.DecodeString(parts[0])
	if err != nil {
		return "", fmt.Errorf("failed to decode key: %v", err)
	}

	// Create AES cipher
	block, err := aes.NewCipher(key)
	if err != nil {
		return "", err
	}

	// Decrypt data using CBC mode
	cbc := cipher.NewCBCDecrypter(block, iv)
	if err != nil {
		return "", err
	}

	// Decode the base64 encoded encrypted data
	payload, err := base64.StdEncoding.DecodeString(encryptedData)
	if err != nil {
		return "", err
	}

	decryptedData := make([]byte, len(payload))
	cbc.CryptBlocks(decryptedData, payload)

	// Remove PKCS7 padding
	decryptedData, err = pkcs7Unpad(decryptedData, aes.BlockSize)
	if err != nil {
		return "", fmt.Errorf("failed to remove padding: %v", err)
	}

	return string(decryptedData), nil
}